// main.cpp : main
// 2009-09

#include "utils.h"
#include "index.h"
#include "record.h"

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <getopt.h>
#include <errno.h>
#include <iconv.h>
#include <assert.h>
#include <string.h>
#include <ctype.h>
#include <sys/socket.h>
#include <sys/epoll.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <netdb.h>

#include <string>
#define PROG_NAME "iplook"

static void print_usage( void );
static int get_data_file_path( std::string &path, const char *input_path = NULL );
static int gbk_to_utf8( std::string &str );
static int get_ip_from_host( uint32_t &ip, bool &is_domain_name, const char *host );
static int find_index_entry( struct index_t &index, uint32_t ip, size_t begin, size_t end );

int main( int argc, char **argv )
{
	bool dump_flag = false;
	bool data_version_flag = false;
	bool origin_format_flag = false;
	const char *input_file_path = NULL;

	struct option long_options[] = {
		{ "file", 1, 0, 'f' },
		{ "dump", 0, 0, 'd' },
		{ "origin-encoding", 0, 0, 'o' },
		{ "data-version", 0, 0, 'V' },
		{ "help", 0, 0, 'h' },
		{ 0, 0, 0, 0 }
	};
	const char *short_options = "f:doVh";
	int opt;

	while( ( opt = getopt_long( argc, argv, short_options, long_options, NULL ) ) != -1 ) {
		switch( opt ) {
		case 'f':
			input_file_path = optarg;
			break;
		case 'd':
			dump_flag = true;
			break;
		case 'o':
			origin_format_flag = true;
			break;
		case 'V':
			data_version_flag = true;
			break;
		case 'h':
		default:
			print_usage();
			exit( EXIT_FAILURE );
		}
	}
	if( optind >= argc && !data_version_flag && !dump_flag ) {
		print_usage();
		exit( EXIT_FAILURE );
	}

	// load file
	FILE *fp = NULL;
	std::string file_path;
	if( ( 0 != get_data_file_path( file_path, input_file_path ) ) ||
			( ( fp = fopen( file_path.c_str(), "r" ) ) == NULL ) ) {
		if( NULL != input_file_path )
			LOG_CRIT( "can't read file : %s", input_file_path );
		LOG_CRIT( "can't read file : \"./qqwry.dat\" or \"~/.share/iplook/qqwry.dat\"" );
	}

	// read head, index, reocrd
	uint32_t index_head = 0, index_end = 0;
	if( fread( &index_head, 4, 1, fp ) != 1 || fread( &index_end, 4, 1, fp ) != 1 )
		LOG_CRIT( "bad data format : %s", file_path.c_str() );

	if( 0 != g_indexs.LoadAllIndex( fp, index_head, ( index_end - index_head ) / 7 + 1 ) )
		LOG_CRIT( "bad data format : %s", file_path.c_str() );
	if( 0 != g_records.Load( fp, index_end ) )
		LOG_CRIT( "bad data format : %s", file_path.c_str() );
	fclose( fp );

	// print data version
	Record record;
	if( data_version_flag ) {
		if( 0 != record.FromAddr( g_indexs.data[g_indexs.size - 1].addr ) )
			LOG_CRIT( "bad data format : %s", file_path.c_str() );
		if( 0 != gbk_to_utf8( record.country ) || 0 != gbk_to_utf8( record.area ) )
			exit( EXIT_FAILURE );
		printf( "%s VERSION :\n%s\t%s\n", file_path.c_str(), record.country.c_str(), record.area.c_str() );
		exit( EXIT_SUCCESS );
	}

	// dump all data
	if( dump_flag ) {
		for( size_t i = 0; i < g_indexs.size - 1; i++ ) {
			PR_IP( g_indexs.data[i].ip );
			record.Reset();
			if( 0 != record.FromAddr( g_indexs.data[i].addr ) )
				LOG_CRIT( "bad data format : %s", file_path.c_str() );
			printf( "\t-\t" );
			PR_IP( record.ip );
			if( !origin_format_flag ) {
				if( 0 != gbk_to_utf8( record.country ) || 0 != gbk_to_utf8( record.area ) )
					exit( EXIT_FAILURE );
			}
			printf( "\t%s\t%s\n", record.country.c_str(), record.area.c_str() );
		}
		return 0;
	}

	// get ip or domain location
	int exit_value = EXIT_FAILURE;
	for(; optind < argc; optind++ ) {
		const char *name = argv[optind];
		uint32_t ip = 0;
		bool is_domain = false;
		if( 0 == strlen( name ) || get_ip_from_host( ip, is_domain, name ) != 0 ) {
			LOG_ERR( "get ip from [%s] failed", name );
			continue;
		}
		/*
		   if( is_domain ) {
			printf( "INFO : GET %s IP ", name );
			PR_IP( ip ); printf( "\n" );
		   }
		 */

		struct index_t index;
		find_index_entry( index, ip, 0, g_indexs.size - 1 );
		record.Reset();
		if( 0 != record.FromAddr( index.addr ) )
			LOG_CRIT( "bad data format : %s", file_path.c_str() );
		if( ip < index.ip || ip > record.ip ) {
			LOG_ERR( "can't find ip [%x], host [%s]", ip, name );
			continue;
		}

		if( !origin_format_flag ) {
			if( 0 != gbk_to_utf8( record.country ) || 0 != gbk_to_utf8( record.area ) )
				exit( EXIT_FAILURE );
		}
		printf( "%s", name );
		if( is_domain ) {
			printf( " [" ); PR_IP( ip ); printf( "]" );
		}
		printf( " FROM %s %s\n", record.country.c_str(), record.area.c_str() );
		exit_value = EXIT_SUCCESS;

	}
	exit( exit_value );
}

static void print_usage( void )
{
	printf( PROG_NAME " : lookup ip location\n"
			"Usage: " PROG_NAME " [OPTIONS]... [IP]... [DOMAIN]...\n"
			"options:\n"
			"    --file, -f file             ip data file.\n"
			"    --dump, -d file             dump out all ip location.\n"
			"    --origin-encoding, -f       print location in origin encoding[GBK].\n"
			"                                it will speed up --dump.\n"
			"    --data-version, -V          print data version.\n"
			"    --help, -h                  print useage.\n" );
}

int get_data_file_path( std::string &path, const char *input_path )
{
	if( NULL != input_path ) {
		if( 0 == access( input_path, R_OK ) ) {
			path = input_path;
			return 0;
		}
		return -1;
	}

	if( 0 == access( "./qqwry.dat", R_OK ) ) {
		path = "./qqwry.dat";
		return 0;
	}

	const char *home = getenv( "HOME" );
	if( NULL != home ) {
		std::string tmp = std::string( home ) + "/.share/iplook/qqwry.dat";
		if( 0 == access( tmp.c_str(), R_OK ) ) {
			path = tmp;
			return 0;
		}
	}

	return -1;
}

static int gbk_to_utf8( std::string &str )
{
	if( str.empty() ) return 0;

	iconv_t ic = iconv_open( "utf8", "gbk" );
	assert( (iconv_t)-1 != ic );

	size_t src_len = str.size();
	size_t dst_len = src_len / 2 * 3 + 4;
	char *src = new char[src_len + 1];
	memcpy( src, str.c_str(), src_len ); src[src_len] = 0;
	char *buf = new char[dst_len];
	memset( buf, 0, dst_len );
	char *src_bak = src, *buf_bak = buf;

	int rc = iconv( ic, &src_bak, &src_len, &buf_bak, &dst_len );
	if( rc < 0 ) {
		LOG_ERR( "convert from gbk to utf8 failed, [%s]", str.c_str() );
	}
	else
		str = buf;

	delete [] buf;
	delete [] src;
	iconv_close( ic );
	return rc < 0 ? -1 : 0;
}

int get_ip_from_host( uint32_t &ip, bool &is_domain, const char *host )
{
	struct in_addr addr;

	memset( &addr, 0, sizeof(addr) );
	if( inet_aton( host, &addr ) == 0 ) {
		is_domain = true;
		char buf[512];
		struct hostent hentry;
		int herrno = 0;
		struct hostent *hp;

		if( gethostbyname_r( host, &hentry, buf, sizeof(buf), &hp, &herrno ) != 0 ) {
			LOG_ERR( "gethostbyname failed, host %s, errno %d", host, errno );
			return -1;
		}
		else {
			memcpy( &addr, hentry.h_addr_list[0], sizeof(struct in_addr) );
		}
	}
	const unsigned char *ip_from = (const unsigned char *)&addr.s_addr;
	unsigned char *ip_to = (unsigned char *)&ip;
	for( int i = 0; i < 4; i++ )
		ip_to[i] = ip_from[3 - i];

	return 0;
}

int find_index_entry( struct index_t &index, uint32_t ip, size_t begin, size_t end )
{
	assert( begin <= end );
	if( begin == end || g_indexs.data[end].ip <= ip ) {
		index = g_indexs.data[end];
		return 0;
	}
	else if( begin + 1 == end ) {
		index = g_indexs.data[begin];
		return 0;
	}
	else {
		size_t m = begin + ( end - begin ) / 2;
		if( g_indexs.data[m].ip < ip )
			return find_index_entry( index, ip, m, end );
		else
			return find_index_entry( index, ip, begin, m );
	}
}

